"""
LeetCode 热题

给你一个整数数组 nums 和一个整数 k ，请你统计并返回 该数组中和为 k 的子数组的个数 。

子数组是数组中元素的连续非空序列。

我们可以通过前缀和的方式来解决这个问题。具体步骤如下：
使用一个字典来记录前缀和出现的次数。
遍历数组 nums，计算当前的前缀和。
检查当前前缀和减去 k 是否在字典中，如果在，则说明存在一个子数组的和为 k，将对应的计数加到结果中。
更新字典中当前前缀和的计数。
"""

from collections import defaultdict

class Solution(object):
    def subarraySum(self, nums, k):
        """
        统计并返回数组中和为 k 的子数组的个数。

        参数:
        nums: 整数数组
        k: 目标和

        返回:
        一个整数，表示和为 k 的子数组的个数

        """
        # 记录前缀和出现的次数
        prefix_sum_count = defaultdict(int)
        # 初始化前缀和为 0 的情况
        prefix_sum_count[0] = 1

        # 当前前缀和
        current_prefix_sum = 0
        # 结果计数
        count = 0

        for num in nums:
            # 更新当前前缀和
            current_prefix_sum += num
            # 检查是否存在当前前缀和减去 k 的前缀和
            if (current_prefix_sum - k) in prefix_sum_count:
                count += prefix_sum_count[current_prefix_sum - k]
            # 更新当前前缀和的计数
            prefix_sum_count[current_prefix_sum] += 1

        return count

if __name__ == '__main__':
    a = Solution()
    # 示例测试
    print(a.subarraySum([1, 1, 1], 2))  # 输出应为 2
